import {toValidRgb, inverseDecorrelate} from './color.js';
import {pixelImage, laplacianPyramid,
makeLaplacianPyramidFromImgData, baseImageLaplacianPyramid} from './spatial.js';
import * as tf from '@tensorflow/tfjs';
/**
* Returns input function and variables with naive pixel space parametrization of an existing image.
* @param {Uint8Array} imgArray Input image data
* @param {*} w width
* @param {*} h height
* @param {*} ch channels
* @param {*} batch batchsize
* @param {*} decorrelate decorrelate colors
*/
export function naiveFromImage(
imgArray, w, h=undefined, ch=3, batch=undefined, decorrelate=true) {
const _w = w;
const _h = h ? h : w;
const _batch = batch ? batch : 1;
const _channels = ch;
const shape = [_batch, _w, _h, _channels];
const imgShape = [_batch, imgArray.height, imgArray.width, 4];
const _decorrelate = _channels === 3 ? decorrelate : false;
let trainable = tf.tidy(() => {
const arrayCopy = Uint8Array.from(imgArray.data);
const origImg = tf.tensor(arrayCopy, imgShape, 'float32');
const resizedRGBA = tf.image.resizeBilinear(origImg, [_w, _h]).div(255);
let wantedChannels = resizedRGBA.slice([0, 0, 0, 0], [_batch, _w, _h, ch]);
if(decorrelate){
wantedChannels = inverseDecorrelate(wantedChannels);
}
return wantedChannels;
});
trainable = [tf.variable(trainable)];
const t = (trainableVars) =>{
let rgb = toValidRgb(trainableVars[0].slice(
[0, 0, 0, 0], [_batch, _w, _h, ch]), _decorrelate, false);
return rgb;
}
return [t, trainable];
}
/**
* Returns input function and variables with naive pixel space parametrization with random initialization.
* @param {*} w width
* @param {*} h height
* @param {*} ch channels
* @param {*} batch batchsize
* @param {*} sd standard deviation
* @param {*} decorrelate decorrelate colors
* @param {*} alpha use alpha channel
*/
export function randImage(
w, h=undefined, ch=3, batch=undefined, sd=undefined,
decorrelate=false, alpha = false
) {
const _w = w;
const _h = h ? h : w;
const _batch = batch ? batch : 1;
const _channels = alpha ? ch+1 : ch;
const shape = [_batch, _w, _h, _channels];
const _decorrelate = _channels === 3 ? decorrelate : false;
let paramF, ret, t, trainable;
paramF = pixelImage;
trainable = [paramF(shape,sd)];
t = (trainableVars) => {
const rgb = toValidRgb(trainableVars[0].slice([0, 0, 0, 0], [_batch, _w, _h, ch]),
_decorrelate, true);
if (alpha) {
const a = tf.sigmoid(trainableVars[0].slice([0, 0, 0, ch], [_batch, _w, _h, 1]));
return tf.concat([rgb, a], -1);
}
return rgb;
}
return [t, trainable];
}
/**
* Returns input function and variables with laplacian pyramid parametrization and random initialization.
* @param {*} w width
* @param {*} h height
* @param {*} ch channels
* @param {*} batch batchsize
* @param {*} sd standard deviation
* @param {*} decorrelate decorrelate colors
* @param {*} nLevels number of pyramid layers
*/
export function randLaplacianPyramid(
w, h=undefined, ch=3, batch=undefined, sd=undefined,
decorrelate=true, nLevels=6
) {
const _w = w;
const _h = h ? h : w;
const _batch = batch ? batch : 1;
const _channels = ch;
const _decorrelate = _channels === 3 ? decorrelate : false;
const shape = [_batch, _w, _h, _channels];
const ret = laplacianPyramid(shape, sd, nLevels);
const f = ret[0];
const trainable = ret[1];
let t = (xArray, weights) => {
return toValidRgb(f(xArray, weights).slice([0, 0, 0, 0], [_batch, _w, _h, ch]),
_decorrelate, false);
}
return [t, trainable];
}
/**
* Returns input function and variables in naive pixel space parametrization with image initialization.
* @param {*} imgArray input image
* @param {*} w width
* @param {*} h height
* @param {*} ch channels
* @param {*} batch batchsize
* @param {*} decorrelate decorrelate colors
* @param {*} nLevels number of pyramid layers
*/
export function imgLaplacianPyramid(imgArray,
w, h=undefined, ch=3, batch=undefined, decorrelate=true, nLevels=6) {
const _w = w;
const _h = h ? h : w;
const _batch = batch ? batch : 1;
const _channels = ch;
const _decorrelate = _channels === 3 ? decorrelate : false;
const shape = [_batch, _w, _h, _channels];
const ret = baseImageLaplacianPyramid(shape, imgArray, nLevels, _decorrelate);
const f = ret[0];
const trainable = ret[1];
let t = (xArray) => {
return toValidRgb(f(xArray).slice([0, 0, 0, 0], [_batch, _w, _h, ch]),
_decorrelate, false);
}
return [t, trainable];
}